package security.crypto.aes;



/*
 *****************************************************************************
 *       AES Decryption.													 *
 *       																	 *
 *       Advanced Encryption Standard implementation in JAVA.   			 *
 *       																	 *
 *       Solution based on Niyaz PK's implementation in C.      			 *
 *		 All credits reserved.												 * 
 *																			 *
 *****************************************************************************
 *   This is the part of the application to decrypt binary files using 		 *
 *   AES algorithm. 														 *
 *****************************************************************************
 * @filename: AES_decrypt.java             	                                 *
 * @language: Java                                                           *
 * @last update at: 2010-04-13                                               *
 *****************************************************************************
 * @students that colaborated on this file:                                  *
 *  55962 - Manuel Monterroso - m.monterroso@gmail.com			             *
 *  57441 - Ant�nio Lages - antoniomslages@ist.utl.pt                        *
 *****************************************************************************
 */
public class AES_decrypt {

	private byte[] rest = new byte[0];
	private final byte[] _initialIV = {(byte)0x19 ,(byte) 0xa5 ,(byte) 0xd6 ,(byte) 0x9d ,(byte) 0xc9 ,(byte) 0xa8 ,(byte) 0xe9 ,(byte) 0xea ,(byte) 0xe6 ,(byte)0x74 ,(byte) 0x1e ,(byte) 0xa5 ,(byte)0x39 ,(byte) 0x35 ,(byte) 0x73 ,(byte) 0xb7};
	private byte[] IV;
	
	private int cmode;
	
	private AES _aes;
	private AES_encrypt _aesEncrypt;
	
	void init(AES aes, AES_encrypt aesEncrypt){
		_aes = aes;
		_aesEncrypt = aesEncrypt;
		
		cmode = aes.getCMode();
		
		IV = new byte[16];		
		System.arraycopy(_initialIV, 0, IV, 0, _initialIV.length);
	}
		
	// This method performs the ciphering of the input data, using the parameters specified in the init method.
	// Note: If the last bytes of the input data do not form a full block, the bytes are not ciphered at that instance.
	byte[] update (byte[] ciphertext) {

		int i, j, round=0;
		// state - the array that holds the intermediate results during decryption
		byte[][] state = new byte[4][4];
		
		int quotient = ciphertext.length / 16; 
		int remainder = ciphertext.length % 16;
		
		if (rest.length > 0) {
			
			byte[] newCiphertext = new byte[ciphertext.length + rest.length];
				
			for(i=0; i < rest.length; i++)
				newCiphertext[i] = rest[i];
			
			for(i=0; i < ciphertext.length; i++)
				newCiphertext[rest.length + i] = ciphertext[i];
			
			quotient = newCiphertext.length / 16;
			remainder = newCiphertext.length % 16;	
			
			ciphertext = newCiphertext;
			
			rest = new byte[0];
		}
		
		if (remainder != 0) {
			
			rest = new byte[remainder];
			
			for(i=0; i < remainder; i++)
				rest[i] = ciphertext[quotient*16 + i];
		}
		
		int ct_length = quotient*16;
		byte[] temp = new byte[16];
		
		for (int globalOffset=0; globalOffset < ct_length;) {
			
			switch(cmode) {
			
			case 2 : { // CTR mode with a constant initialization vector as nounce
				
				//Copy the input IV (as nounce) to state array			
				for(i=0; i < 4; i++) {
					
					for(j=0;j<4;j++) {
						state[j][i] = IV[i*4 + j];
					}
				}
				
				// Add the first round key to the state before starting the rounds
				state = _aesEncrypt.addRoundKey(state, _aes.getRoundKey(), 0);
				
				// There will be Nr rounds.
				// The first Nr-1 rounds are identical.
				// These Nr-1 rounds are executed in the loop below.
				for(round=1; round < _aes.getNr(); round++) {
					
					state = _aesEncrypt.addRoundKey(_aesEncrypt.mixColumns(_aesEncrypt.shiftRows(_aesEncrypt.subBytes(state))), _aes.getRoundKey(), round);
				}
				
				// The last round is given below
				// The MixColumns function is not here in the last round
				state = _aesEncrypt.addRoundKey(_aesEncrypt.shiftRows(_aesEncrypt.subBytes(state)), _aes.getRoundKey(), round);

				// The encryption process is over. Copy the state array to output array.
				for(i=0;i<4;i++) {
					
					for(j=0;j<4;j++) {
						ciphertext[globalOffset + i*4 + j] ^= state[j][i];
					}
				}
				
				break; 
			}
			
			default : { // ECB or CBC modes
				
				if (cmode==1) // CBC mode
					for (i=0; i < 16; i++)
						temp[i] = ciphertext[globalOffset + i];
				
				//Copy the input ciphertext to state array
				for(i=0; i < 4; i++) {
					
					for(j=0;j<4;j++) {
						state[j][i] = ciphertext[globalOffset + i*4 + j];
					}
				}
				
				// Add the First round key to the state before starting the rounds.
				state = addRoundKey(state, _aes.getRoundKey(), _aes.getNr()); 

				for(round=((_aes.getNr())-1); round > 0; round--){
					
					state = invMixColumns(addRoundKey(invSubBytes(invShiftRows(state)), _aes.getRoundKey(), round));
				}	
				
				state = addRoundKey(invSubBytes(invShiftRows(state)), _aes.getRoundKey(),0);

				for(i=0;i<4;i++) {
					
					for(j=0;j<4;j++) {
						ciphertext[globalOffset + i*4 + j] = state[j][i];
					}
				}
				
				if (cmode==1) { // CBC mode
					for (i=0; i < 16; i++)
						ciphertext[globalOffset + i] ^= IV[i];
					
					for (i=0; i < 16; i++)
						IV[i] = temp[i];
				}
				
			}
			}
			
			globalOffset+=16;
		}
		
		return ciphertext;
	}

	// This method performs the ciphering of the input data, using the parameters specified in the init method.
	// Note: If the last bytes of the input data do not form a full block, the bytes are not ciphered at that instance.
	byte[] update (byte[] ciphertext, int inputlen) {
		
		int i, j, round=0;
		// state - the array that holds the intermediate results during decryption
		byte[][] state = new byte[4][4];
		
		int quotient = inputlen / 16; 
		int remainder = inputlen % 16;
		
		if (rest.length > 0) {
			
			byte[] newCiphertext = new byte[inputlen + rest.length];
				
			for(i=0; i < rest.length; i++)
				newCiphertext[i] = rest[i];
			
			for(i=0; i < inputlen; i++)
				newCiphertext[rest.length + i] = ciphertext[i];
			
			quotient = newCiphertext.length / 16;
			remainder = newCiphertext.length % 16;	
			
			ciphertext = newCiphertext;
			
			rest = new byte[0];
		}
		
		if (remainder != 0) {
			
			rest = new byte[remainder];
			
			for(i=0; i < remainder; i++)
				rest[i] = ciphertext[quotient*16 + i];
		}
		
		int ct_length = quotient*16;
		byte[] temp = new byte[16];
		
		for (int globalOffset=0; globalOffset < ct_length;) {
			
			switch(cmode) {
			
			case 2 : { // CTR mode with a constant initialization vector as nounce
				
				//Copy the input IV (as nounce) to state array			
				for(i=0; i < 4; i++) {
					
					for(j=0;j<4;j++) {
						state[j][i] = IV[i*4 + j];
					}
				}
				
				// Add the first round key to the state before starting the rounds
				state = _aesEncrypt.addRoundKey(state, _aes.getRoundKey(), 0);
				
				// There will be Nr rounds.
				// The first Nr-1 rounds are identical.
				// These Nr-1 rounds are executed in the loop below.
				for(round=1; round < _aes.getNr(); round++) {
					
					state = _aesEncrypt.addRoundKey(_aesEncrypt.mixColumns(_aesEncrypt.shiftRows(_aesEncrypt.subBytes(state))), _aes.getRoundKey(), round);
				}
				
				// The last round is given below
				// The MixColumns function is not here in the last round
				state = _aesEncrypt.addRoundKey(_aesEncrypt.shiftRows(_aesEncrypt.subBytes(state)), _aes.getRoundKey(), round);

				// The encryption process is over. Copy the state array to output array.
				for(i=0;i<4;i++) {
					
					for(j=0;j<4;j++) {
						ciphertext[globalOffset + i*4 + j] ^= state[j][i];
					}
				}
				
				break; 
			}
			
			default : { // ECB or CBC modes
				
				if (cmode==1) // CBC mode
					for (i=0; i < 16; i++)
						temp[i] = ciphertext[globalOffset + i];
				
				//Copy the input ciphertext to state array
				for(i=0; i < 4; i++) {
					
					for(j=0;j<4;j++) {
						state[j][i] = ciphertext[globalOffset + i*4 + j];
					}
				}
				
				// Add the First round key to the state before starting the rounds.
				state = addRoundKey(state, _aes.getRoundKey(), _aes.getNr()); 

				for(round=((_aes.getNr())-1); round > 0; round--){
					
					state = invMixColumns(addRoundKey(invSubBytes(invShiftRows(state)), _aes.getRoundKey(), round));
				}	
				
				state = addRoundKey(invSubBytes(invShiftRows(state)), _aes.getRoundKey(),0);

				for(i=0;i<4;i++) {
					
					for(j=0;j<4;j++) {
						ciphertext[globalOffset + i*4 + j] = state[j][i];
					}
				}
				
				if (cmode==1) { // CBC mode
					for (i=0; i < 16; i++)
						ciphertext[globalOffset + i] ^= IV[i];
					
					for (i=0; i < 16; i++)
						IV[i] = temp[i];
				}
				
			}
			}
			
			globalOffset+=16;
		}
		
		return ciphertext;
	}
	
	// This method performs the ciphering of the input data, using the parameters specified in the init
	// method. This method also performs the necessary procedures to conclude the ciphering of the
	// data stream, including the addition of necessary padding.
	byte[] doFinal (byte[] ciphertext) {
		
		int i; 
		
		int quotient = ciphertext.length / 16; 
		int remainder = ciphertext.length % 16;
		
		if (rest.length > 0) {
			
			byte[] newCiphertext = new byte[ciphertext.length + rest.length];
				
			for(i=0; i < rest.length; i++)
				newCiphertext[i] = rest[i];
			
			for(i=0; i < ciphertext.length; i++)
				newCiphertext[rest.length + i] = ciphertext[i];
			
			quotient = newCiphertext.length / 16;
			remainder = newCiphertext.length % 16;	
			
			ciphertext = newCiphertext;
			
			rest = new byte[0];
		}
		
		if(remainder != 0) {
			System.out.println("Something is wrong with doFinal function in decrypt!");
			System.exit(-1);
		}
			
		ciphertext = update(ciphertext);
		
		byte[] block = new byte[16];
		
		for(i=0; i < 16; i++)
			block[i] = ciphertext[(quotient-1)*16 + i];

		block = invPadding(block);
		
		byte[] newCiphertext = new byte[(quotient-1)*16 + block.length];
		
		for(i=0; i < (quotient-1)*16; i++)
			newCiphertext[i] = ciphertext[i];
		
		for(i=0; i < block.length; i++)
			newCiphertext[(quotient-1)*16 + i] = block[i];

		return newCiphertext;
	}
	
	// This method performs the ciphering of the input data, using the parameters specified in the init
	// method. This method also performs the necessary procedures to conclude the ciphering of the
	// data stream, including the addition of necessary padding.	
	byte[] doFinal (byte[] ciphertext, int inputlen) {
		
		int i; 
		
		int quotient = inputlen / 16; 
		int remainder = inputlen % 16;
		
		if (rest.length > 0) {
			
			byte[] newCiphertext = new byte[inputlen + rest.length];
				
			for(i=0; i < rest.length; i++)
				newCiphertext[i] = rest[i];
			
			for(i=0; i < inputlen; i++)
				newCiphertext[rest.length + i] = ciphertext[i];
			
			quotient = newCiphertext.length / 16;
			remainder = newCiphertext.length % 16;	
			
			ciphertext = newCiphertext;
			
			rest = new byte[0];
		}
		
		if(remainder != 0) {
			System.out.println("Something is wrong with doFinal function in decrypt!");
			System.exit(-1);
		}
			
		ciphertext = update(ciphertext, quotient*16);
		
		byte[] block = new byte[16];
		
		for(i=0; i < 16; i++)
			block[i] = ciphertext[(quotient-1)*16 + i];

		block = invPadding(block);
		
		byte[] newCiphertext = new byte[(quotient-1)*16 + block.length];
		
		for(i=0; i < (quotient-1)*16; i++)
			newCiphertext[i] = ciphertext[i];
		
		for(i=0; i < block.length; i++)
			newCiphertext[(quotient-1)*16 + i] = block[i];
		
		return newCiphertext;
	}
		
	//***************************************************************************************************
	// This funtion does the inverse process correspondent to the padding function 
	// to aquire the original block before encryption
	byte[] invPadding (byte[] block) {
		
		int length = block.length;
		
		if(length != 16) {
			System.out.println("Padding block must have 128 bits (16 bytes)!");
			System.exit(-1);
		}
		
		int n = (int) block[15];
		
		byte[] result = new byte[16-n];
		
		for(int i=0; i < result.length; i++)
			result[i] = block[i];
		
		return result;
	}	

	// This function adds the round key to state.
	// The round key is added to the state by an XOR function.
	byte[][] addRoundKey(byte[][] state, byte[] roundKey, int round) {
		
		int i,j;
		
		for(i=0;i<4;i++) {
		
			for(j=0;j<4;j++) {
				
				state[j][i] ^= roundKey[round * _aes.getNb() * 4 + i * _aes.getNb() + j];
			}
		}
		
		return state;
	}
	
	// MixColumns function mixes the columns of the state matrix
	byte[][] invMixColumns(byte[][] state) {
		
		int i;
		byte a,b,c,d;
		
		for(i=0;i<4;i++) {	
			
			a = state[0][i];
			b = state[1][i];
			c = state[2][i];
			d = state[3][i];

			state[0][i] = (byte) (multiply(a, (byte)0x0e) ^ multiply(b, (byte)0x0b) ^ multiply(c,(byte)0x0d) ^ multiply(d, (byte)0x09));
			state[1][i] = (byte) (multiply(a, (byte)0x09) ^ multiply(b, (byte)0x0e) ^ multiply(c,(byte)0x0b) ^ multiply(d, (byte)0x0d));
			state[2][i] = (byte) (multiply(a, (byte)0x0d) ^ multiply(b, (byte)0x09) ^ multiply(c,(byte)0x0e) ^ multiply(d, (byte)0x0b));
			state[3][i] = (byte) (multiply(a, (byte)0x0b) ^ multiply(b, (byte)0x0d) ^ multiply(c,(byte)0x09) ^ multiply(d, (byte)0x0e));
		}
		
		return state;
	}
	
	// The invShiftRows() function shifts the rows in the state to the left.
	// Each row is shifted with different offset.
	// Offset = Row number. So the first row is not shifted.
	byte[][] invShiftRows(byte[][] state){
		
		byte temp;
		
		temp=state[1][3];// Rotate first row 1 columns to right	
		temp=state[1][3];
		state[1][3]=state[1][2];
		state[1][2]=state[1][1];
		state[1][1]=state[1][0];
		state[1][0]=temp;

		// Rotate second row 2 columns to right	
		temp=state[2][0];
		state[2][0]=state[2][2];
		state[2][2]=temp;

		temp=state[2][1];
		state[2][1]=state[2][3];
		state[2][3]=temp;

		// Rotate third row 3 columns to right
		temp=state[3][0];
		state[3][0]=state[3][1];
		state[3][1]=state[3][2];
		state[3][2]=state[3][3];
		state[3][3]=temp;

		return state;
	}
	
	// The SubBytes Function Substitutes the values in the
	// state matrix with values in an S-box.
	byte[][] invSubBytes(byte[][] state) {
		
		int i,j;
		
		for(i=0;i<4;i++) {
			
			for(j=0;j<4;j++) {
				
				
				state[i][j] = (byte) _aes.getSBoxInvertValue(state[i][j]);
			}
		}
		
		return state;
	}
	
	// xtime is a function that finds the product of {02} and the argument to xtime modulo {1b}  
	byte xtime (byte x) {
		
		return (byte) ((x<<1) ^ (((x>>7) & 1) * 0x1b));
	}
	
	// Multiplty is a function used to multiply numbers in the field GF(2^8)
	byte multiply(byte x, byte y){
		
		return (byte) (((y & 1) * x) ^ ((y>>1 & 1) * xtime(x)) ^ ((y>>2 & 1) * xtime(xtime(x))) ^ ((y>>3 & 1) * xtime(xtime(xtime(x)))) ^ ((y>>4 & 1) * xtime(xtime(xtime(xtime(x))))));
	}
}
